﻿#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>  
#include <direct.h> 
#include <vector>
#include <thread>
#include <chrono>

#include "model_trans_seg_v2.h"
#include "../public/file_utility.hpp"
#include "../public/gpu_utility.hpp"

void enum_models(std::string& mon_folder, std::string& dest_folder,std::string & mark_info)
{
	namespace fs = boost::filesystem;
	
	std::string gpu_model = GetGPUModel();
	if (gpu_model == "")
	{
		std::cout << "Unknown Gpu Model! \n";
		return;
	}
	//检查目标文件夹是否有效
	fs::path _dest_folder(dest_folder);
	if (!fs::exists(_dest_folder))
	{
		return;
	}
	std::vector<std::string> cfg_files = EnumFiles(mon_folder, std::string(".yaml"));
	for (int i = 0; i < cfg_files.size(); ++i)
	{
		st_trans_model_input input;
		fs::path _cfg_path(cfg_files[i]);
		input.cfg_path = _cfg_path.string();

		fs::path _model_path; 
		{
			fs::path _path = _cfg_path.replace_extension(".ennx");
			if (fs::exists(_path))
			{
				_model_path = _path;
			}
		} 
		{
			fs::path _path = _cfg_path.replace_extension(".onnx");
			if (fs::exists(_path))
			{
				_model_path = _path;
			}
		}
		if (_model_path.empty())
		{
			continue;
		}
		
		fs::path dats_path(dest_folder + "\\" + _cfg_path.stem().string() + mark_info + ".seg2");
		if (fs::exists(dats_path))
		{
			continue;
		}

		std::cout << "START===============================================\n";
		input.model_path = _model_path.string();
		input.gpu_model = gpu_model;
		input.trans_path = dats_path.string() + "-tmp";

		ModelTransSegV2 c;
		if (c.TransModel(input))
		{
			std::cout << "DONE================================================\n";
			int result = std::rename(input.trans_path.c_str(), dats_path.string().c_str());
			if (result == 0) 
			{
				std::cout << "文件重命名成功。\n";
			} 
			else 
			{
				std::cout << "文件重命名失败。\n";
			}
		}
		else
		{
			std::cout << "FAIL!================================================\n";
		}
	}
}

int main(int argc, char* argv[])
{
	try
	{
		char cur_folder[260];
		_getcwd(cur_folder, sizeof(cur_folder));
		//printf("current working directory: \n %s\n", cur_folder);
		std::string cwd = cur_folder;

		std::string mon_folder = cwd;
		std::string dest_folder = cwd;
		std::string mark_info = "";
		if (argc >= 2)
		{
			mon_folder = argv[1];
			dest_folder = mon_folder;
		}
		if (argc >= 3)
		{
			dest_folder = argv[2];
		}
		if (argc >= 4)
		{
			mark_info = argv[3];
		}

		enum_models(mon_folder, dest_folder, mark_info);
		printf("done!\n");

	}
	catch (...)
	{
		printf("exception !\n");
	}

	for (int i = 0; i<5; ++i)
	{
		std::this_thread::sleep_for(std::chrono::milliseconds(1000));
		//printf("wait for shutdown %d !\n", 5 - i);
	}

	//getchar();
}
